#!/usr/bin/env python
#
# Wakeonlan daemon to keep certain machines awake, this has now been
# made dynamic by inserting and monitoring iptables data
#

import re

#
# Config
#
MACHINES = {
  '192.168.0.1'   : [ '00:1a:4d:42:49:ee', '00:40:f4:4a:09:41' ],
  '192.168.0.2'   : [ '00:1a:4d:42:49:ee', '00:40:f4:4a:09:41' ],
  '192.168.0.255' : [ '00:1a:4d:42:49:ee', '00:40:f4:4a:09:41' ],
  # TODO: would be useful if this could be port specific
}
PORT     = 9
REPEAT   = 2
LOG_PATH = '/tmp/wakeonlan.log'

# ###########################################################################
# Logging
# ###########################################################################

def log ( msg ):
  from datetime import datetime
  tm  = datetime.now().strftime('%F %T')
  msg = '%s - %s' % (tm, msg)
  try:
    open(LOG_PATH, 'a').write(msg + '\n')
  except: pass
  print msg

# ###########################################################################
# IPtables
# ###########################################################################

# Convert number
nume = re.compile('(\d+)([KMG]?)')
def number ( num ):
  import re
  ret = 0
  r   = nume.search(num)
  if r:
    ret = int(r.group(1))
    if   r.group(2) == 'K': ret *= 1000
    elif r.group(2) == 'M': ret *= 1000000
    elif r.group(2) == 'G': ret *= 1000000000
  return ret

# Get current rules
def iptables_list_rules ():
  import string
  from subprocess import Popen, PIPE

  # Execute
  proc = Popen(['/usr/sbin/iptables-save'], stdout=PIPE)
  ret  = map(string.strip, proc.stdout.readlines())
  
  # Return
  return ret
      
# Add rules
def iptables_add_rules ( ips ):
  from subprocess import call

  # Get list
  rules = iptables_list_rules()
  log('iptables rules:')
  for r in rules:
    log('  ' + r)

  # Add missing rules (at start)
  for ip in ips:
    rule = '-A OUTPUT -d %s/32' % ip
    if rule not in rules:
      log('add %s' % rule)
      call('iptables -I OUTPUT 1 -d %s' % ip, shell=True)
    rule = '-A INPUT -s %s/32' % ip
    if rule not in rules:
      log('add %s' % rule)
      call('iptables -I INPUT 1 -s %s' % ip, shell=True)

# Get statistics
def iptables_get_stats ( ips ):
  import re
  from subprocess import Popen, PIPE
  ret   = {}
  chain = None

  # Initialise
  for ip in ips:
    ret[ip] = { 'rx' : 0, 'tx' : 0 }

  # Execute
  proc = Popen(['iptables', '-vnL'], stdout=PIPE)

  # Process
  for l in proc.stdout.readlines():
    if l.startswith('Chain'):
      chain = l.split(' ')[1]
    elif chain in [ 'OUTPUT', 'INPUT' ]:
      pts = re.split('\s+', l.strip())
      if len(pts) < 8: continue
      if chain == 'OUTPUT':
        d  = 'tx'
        ip = pts[7]
      else:
        d  = 'rx'
        ip = pts[6]
      if ip in ips:
        ret[ip][d] += number(pts[0])
    
  # Done
  return ret
  
# ###########################################################################
# Device
# ###########################################################################

import ctypes
__voidptrsize = ctypes.sizeof(ctypes.c_voidp)
_64bit = (__voidptrsize == 8)
_32bit = (__voidptrsize == 4)
if _64bit:
  INPUTEVENT_STRUCT = "=LLLLHHl"
  INPUTEVENT_STRUCT_SIZE = 24
elif _32bit: # 32bit
  INPUTEVENT_STRUCT = "=iiHHi"
  INPUTEVENT_STRUCT_SIZE = 16
else:
  raise RuntimeError("Couldn't determine architecture, modify " + __file__ +
                     " to support your system.")

def input_setup ( path = '/dev/input/event2' ):
  import select
  pd = select.poll()
  fd = open(path)
  pd.register(fd, select.POLLPRI | select.POLLIN)
  return (pd, fd)
  
def input_ready ( pd, timeout = 0.1 ):
  import select, os
  for (fd,e) in pd.poll(timeout):
    if e & (select.POLLIN | select.POLLPRI):
      os.read(fd, INPUTEVENT_STRUCT_SIZE)
      return True
    else:
      print 'ERROR: unhandled poll event %d' % e
  return False

# ###########################################################################
# MAC
# ###########################################################################

# MAC expression
mace = re.compile('^[0-9A-F]{2}(:[0-9A-F]{2}){5}$')

# Check if mac
def is_mac ( mac ):
  return mace.search(mac.upper()) is not None

# ###########################################################################
# WakeOnLAN
# ###########################################################################

def send ( dst, port, repeat = 1 ):
  import socket
  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
  mac = ''.join(map(lambda x: chr(int(x, 16)), dst.split(':')))
  data = '\xff' * 6 + mac * 16
  for i in range(repeat): s.sendto(data, ('255.255.255.255', port))

# ###########################################################################
# Main
# ###########################################################################

if __name__ == '__main__':
  import sys
  from time import sleep
  from datetime import datetime
  from optparse import OptionParser

  # Config
  ips  = MACHINES.keys()
  
  # Process comamnd line
  optp = OptionParser()
  optp.add_option('-p', '--period', default=0.1, type='float',
                  help='Define polling period')
  optp.add_option('-n', '--iptables', default=False, action='store_true',
                  help='Monitor iptables for data')
  optp.add_option('-i', '--input', default=False, action='store_true',
                  help='Monitor /dev/input')
  optp.add_option('-d', '--device', default='/dev/input/event2',
                  help='Specify /dev/input event to monitor')
  (opts,args) = optp.parse_args()
                  
  # Setup iptables
  if opts.iptables:
    iptables_add_rules(ips)
    ip_prev = iptables_get_stats(ips)
    ip_pend = []
    
  # Setup /dev/input
  if opts.input:
    (input_pd, input_fd) = input_setup(opts.device)
    
  # Error
  if not opts.iptables and not opts.input:
    log('ERROR: no monitoring specified')
    sys.exit(1)
    
  # Loop forever
  while True:
    try:
    
      # Wake set
      wake = set()
    
      # Wait for input
      if opts.input:
        if input_ready(input_pd, timeout=opts.period):
          for ip in MACHINES:
            for m in MACHINES[ip]:
              wake.add(m)
             
      # Wait period
      else:
        time.sleep(opts.period)
        
      # IPtables stats
      if opts.iptables:
        norx  = []
        hastx = []

        # Get new stats
        stats = iptables_get_stats(ips)
        for ip in ips:
          if stats[ip]['rx'] == ip_prev[ip]['rx']:
            norx.append(ip)
            if stats[ip]['tx'] != ip_prev[ip]['tx']:
              hastx.append(ip)

        # Wake set
        for ip in ip_pend:
          if ip in norx:
            for m in MACHINES[ip]:
              wake.add(m)
  
        # Pending
        ip_pend = []
        for ip in norx:
          if ip in hastx: 
            log('no rx from %s' % ip)
            ip_pend.append(ip)

        # Save stats
        ip_prev  = stats

      # Wakeup
      for m in wake: send(m, PORT, REPEAT)
      
    # Ignore
    except KeyboardInterrupt, e:
      break
    
    # Error
    except Exception, e:
      log('ERROR: %s' % e)
